package org.chorus.utils;

//  *****************************************************************************************************
/** Class for finding the solution of a set of linear algebraic equations via LU decomposition.
  * Adapted from the algorithm documented in Numerical Recipes in C (Press <i>et al</i>, 1988, pp.43-45).
  * @author Jo Wood
  * @version 2.3.1, 17th July 2012.
  */
//  *****************************************************************************************************

/* This file is part of the Chorus river bed analysis software package. Chorus is free software:
 * you can redistribute it and/or modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation, either version 3 of the License, or (at your
 * option) any later version.
 * 
 * Chorus is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
 * without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  
 * See the GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License along with this
 * source code (see COPYING included with this source code). If not, see 
 * http://www.gnu.org/licenses/.
 */

public class LUDecomp
{
    // ----------------- Object variables ------------------
            
    private int[]      index;  // Row permutation from pivoting.  
    private double[][] lud;    // LU decomposed matrix (LHS of equations).
    private double[]   b;      // Vector representing normal e.q. RHS and solved unknowns.
    
    // ------------------- Constructor --------------------
    
    /** Solves a set of linear equations provided by a matrix holding the left hand side and a
     *  vector holding the right hand side. The number of rows corresponds to the number of 
     *  equations, the number of columns corresponds to the number of unknowns. It therefore
     *  follows that the left hand side of the equation must have at least as many rows as columns.
     *  This constructor will convert the expression into a set of normal equations before using 
     *  LU decomposition to solve the equations.
     *  @param lhs Matrix holding LHS of equations to solve.
     *  @param rhs Vector holding RHS of equations to solve.  
     */
    public LUDecomp(double[][] lhs, double[] rhs)
    {
        // Check matrix and vector have same number of rows.
        if (lhs.length != rhs.length)
        {
            System.err.println("Warning: Ordinary Least Squares cannot be applied as matrix and vector are of different sizes.");
            return;
        }
        
        // Create normal equations.      
        double[][] ATdotA = Matrix.mmMultiply(Matrix.mTranspose(lhs),lhs);
        b = Matrix.mvMultiply(Matrix.mTranspose(lhs),rhs);

        // Create LU decomposed matrix.
        createLUMatrix(ATdotA);
        
        if (lud == null)  // Singular matrix.
        {
            b = null;
            return;
        }
        
        // Solve equations.
        luBackSub(b);
    }
    
    /** Creates the LU Decomposed matrix based on LHS of a set of normal equations equations.
      * @param lhs Matrix holding LHS of normal equations to solve.
      */
    public LUDecomp(double lhs[][])
    {
        // Creates an LU decomposed matrix storing the result in lud[][].
        createLUMatrix(lhs);
    }
    
    // -------------------- Methods -----------------------

    /** Reports the solved coefficients. The RHS of the equation must have been given through a
     *  constructor. Alternatively use luBackSub() to calculate on the fly.
     *  @return Solved coefficients.  
     */
    public double[] getCoeff()
    {
        return b;
    }
    
    /** Performs forward and back substitution on an LU decomposed matrix.
     *  @param rhs RHS of equations to be solved. The array of supplied
     *  values gets replaced with the solved coefficients.
     */ 
    public void luBackSub(double rhs[])
    {
        int    row,col,
               ii=-1, ip, n;
        double sum;
        
        n = rhs.length;
        
        for (row=0; row<n; row++)
        {
            ip=index[row];
            sum = rhs[ip];
            rhs[ip] = rhs[row];
            
            if (ii >= 0)
            {
            	for (col=ii; col<row; col++)
            	{
                    sum -= lud[row][col]*rhs[col];
            	}
            }
            else if (sum != 0F)
            {
            	ii = row;
            }         
            rhs[row] = sum;
        }
        
        for (row=n-1; row>=0; row--)
        {
            sum = rhs[row];
            
            for (col=row+1; col<n; col++)
            {
            	sum -= lud[row][col]*rhs[col];
            }
            rhs[row]=sum/lud[row][row];
        }
    }
    
    // ----------------------- Private Methods -------------------------
    
    /** Creates the LU Decomposed matrix based on LHS of a set of normal equations equations.
      * @param lhs Matrix holding LHS of normal equations to solve.
      */
    private void createLUMatrix(double lhs[][])
    {
        int    row,col,
               rowMax,k,n;
    
        double  big, dum,    // Temporary variables.
                sum, temp,
                vv[];        // Implicit scaling of each row.    
                
        n       = lhs.length ;
        vv      = new double[n];
        index   = new int[n];
        lud     = new double[n][n];
        rowMax  = 0;
    
        // Make a copy of matrix to be decomposed.
        for (row=0; row<n; row++)
        {
            for (col=0; col<n; col++)
            {
                lud[row][col] = lhs[row][col];    
            }
        }
                    
        // Find largest value in each matrix row. 
        // This is used for scaling other values.   
        for (row=0; row<n; row++)
        {
            big = 0F;
        
            for (col=0; col<n; col++)
            {
                temp = Math.abs(lud[row][col]);
                if (temp > big)
                {
                    big = temp;
                }
            }
        
            if (big == 0)
            {
                System.err.println("Error: Singular matrix in LUDecomp");
                lud = null;
                return;
            }
        
            vv[row] = 1F/big;
        }
    
        // Process column by column.
        for (col=0; col<n; col++)
        {
            for (row=0; row<col; row++)
            {
                sum = lud[row][col];
            
                for (k=0; k<row; k++)
                {
                    sum -= lud[row][k]*lud[k][col];
                }
                lud[row][col] = sum;
            }
        
            big = 0F;   // Initialse search for largest pivot element.
        
            for (row=col; row<n; row++)
            {
                sum = lud[row][col];
            
                for (k=0; k<col; k++)
                {
                    sum -= lud[row][k]*lud[k][col];
                }
            
                lud[row][col] = sum;
            
                dum = vv[row]*Math.abs(sum);
                if (dum >= big)
                {
                    big = dum;
                    rowMax = row;
                }
            }
        
            if (col != rowMax)
            {
                for (k=0; k<n; k++)
                {
                    dum = lud[rowMax][k];
                    lud[rowMax][k] = lud[col][k];
                    lud[col][k] = dum;
                }
                vv[rowMax] = vv[col];
            }
        
            index[col] = rowMax;
                        
            if (col != (n-1))
            {
                dum = 1F/lud[col][col];
            
                for (row=col+1; row<n; row++)
                {
                    lud[row][col] *= dum;
                }
            }
        }   
    }
}    